import asyncio
import json

from app.api.generate.utils.generate_queue.queue_utils import send_queue_message
from app.db.session import SessionLocal
from app.models.generate import GenerateImageRecord, GenerateImageResult
from app.api.generate.schemas import generate_image_schema

from app.utils.curd_base import ModelCRUD
from app.utils import security
from app.utils.custom_response import success_response
from fastapi import APIRouter, Query, Body, Path, UploadFile
from fastapi import Depends
from sqlalchemy.orm import Session
from app.db.database import get_db
from app.models.system import Users
from typing import Any
from fastapi.encoders import jsonable_encoder
from app.utils.custom_exc import CustomException

from app.utils.file_storage import get_storage_file, delete_storage_file, save_storage_file


topic = "request_generate_image"


class GenerateImageRecordCRUD(ModelCRUD):
    model = GenerateImageRecord

    @staticmethod
    def serializer(queryset, only: tuple = (), date_format: str = None, datetime_format: str = None,
                   time_format: str = None):
        li = []
        for instance in queryset:
            dit = instance.to_dict(only=only, date_format=date_format, datetime_format=datetime_format, time_format=time_format)
            dit['user_nickname'] = instance.user.nickname if instance.user else None
            result_li = []
            for ins in instance.generate_image_result:
                result_dit = ins.to_dict(only=("id", "path", "record_id"), date_format=date_format, datetime_format=datetime_format, time_format=time_format)
                result_dit['path'] = get_storage_file(result_dit['path'])
                result_li.append(result_dit)
            dit['result'] = result_li
            # print(instance.user.id)
            li.append(dit)
        return li
        # return [instance.to_dict(only=only, date_format=date_format, datetime_format=datetime_format,
        #                          time_format=time_format) for instance in queryset]


router = APIRouter()


@router.post("/generate_image", summary="生成图片", response_model=generate_image_schema.GenerateImageRecordSchemaRes, name='生成图片')
async def generate_image(
    *,
    db: Session = Depends(get_db),
    req_data: generate_image_schema.GenerateImageRecordSchemaReq,
    current_user: Users = Depends(security.verify_access_token),
) -> Any:
    crud = GenerateImageRecordCRUD()
    req_data = jsonable_encoder(req_data)
    batch_count = req_data.get('batch_count', 1)

    valid = generate_image_schema.GenerateImageRecordSchemaCreate.validate
    datetime_format = '%Y-%m-%d %H:%M:%S'
    if isinstance(req_data.get('params'), list):
        li = []
        params_li = req_data.get('params')
        for params in params_li:
            # try:
            #     params = json.loads(item)
            # except:
            #     raise CustomException("请提交正确字段")
            req_data['number'] = params.get("number", 1)
            req_data['user_id'] = current_user.id
            req_data['params'] = json.dumps(params, ensure_ascii=False)
            item = crud.create(db, data=valid(req_data))
            item_dict = item.to_dict(only=('id', 'params', 'number', 'type', 'state', 'create_time'), datetime_format=datetime_format)
            item_dict['result'] = []
            item_dict['user_nickname'] = item.user.nickname if item.user else None
            li.append(item_dict)
            # ----------------------------------
            send_dict = {'record_id': item.id, 'params': params}
            send_queue_message(topic, send_dict)
        return success_response(data=li, message="批量生成中...")
    else:

        if batch_count > 1:
            params = req_data.get("params")
            # try:
            #     params = json.loads(req_data.get("params"))
            # except:
            #     raise CustomException("请提交正确字段")
            req_data['number'] = params.get("number", 1)
            req_data['user_id'] = current_user.id
            req_data['params'] = json.dumps(params, ensure_ascii=False)
            li = []
            for i in range(batch_count):
                item = crud.create(db, data=valid(req_data))
                item_dict = item.to_dict(only=('id', 'params', 'number', 'type', 'state', 'create_time'), datetime_format=datetime_format)
                item_dict['result'] = []
                item_dict['user_nickname'] = item.user.nickname if item.user else None
                li.append(item_dict)
                # ----------------------------------
                send_dict = {'record_id': item.id, 'params': params}
                send_queue_message(topic, send_dict)

            return success_response(data=li, message="批量生成中...")
        else:
            params = req_data.get("params")
            # try:
            #     params = json.loads(req_data.get("params"))
            # except:
            #     raise CustomException("请提交正确字段")
            req_data['number'] = params.get("number", 1)
            req_data['user_id'] = current_user.id
            req_data['params'] = json.dumps(params, ensure_ascii=False)
            item = crud.create(db, data=valid(req_data))
            item_dict = item.to_dict(only=('id', 'params', 'number', 'type', 'state', 'create_time'), datetime_format=datetime_format)
            item_dict['result'] = []
            item_dict['user_nickname'] = item.user.nickname if item.user else None
            resp_data = item_dict

            # -------------------------------
            send_dict = {'record_id': item.id, 'params': params}
            send_queue_message(topic, send_dict)

        return success_response(data=resp_data)


@router.get("/get_generate_image_record", summary="获取生成图片记录", response_model=generate_image_schema.GenerateImageRecordSchemaRes, name='获取生成图片记录')
async def get_generate_image(
    *,
    db: Session = Depends(get_db),
    page: int = Query(1, description='页码', ge=1),
    page_size: int = Query(10, description='页面大小', le=50, ge=1),
    current_user: Users = Depends(security.verify_access_token),
) -> Any:
    crud = GenerateImageRecordCRUD()
    queryset, total = crud.get_page_queryset(db, page=page, page_size=page_size)
    datetime_format = '%Y-%m-%d %H:%M:%S'
    items = crud.serializer(queryset, only=('id', 'user_id', 'params', 'number', 'type', 'state', 'create_time'), datetime_format=datetime_format)
    resp_data = {
        "items": items,
        'total': total
    }
    return success_response(data=resp_data)


@router.post("/del_generate_image_record", summary="删除生成图片记录", response_model=generate_image_schema.GenerateImageRecordSchemaRes, name='删除生成图片记录')
async def del_generate_image(
    *,
    db: Session = Depends(get_db),
    req_data: generate_image_schema.GenerateImageRecordSchemaDel,
    current_user: Users = Depends(security.verify_admin_permission),
) -> Any:
    crud = GenerateImageRecordCRUD()
    queryset = crud.get_all_queryset(db)
    req_data = jsonable_encoder(req_data)
    queryset = queryset.filter(crud.model.id.in_(req_data.get('id', [])))
    for instance in queryset:
        result = instance.generate_image_result
        for ins in result:
            delete_storage_file(ins.path)
    queryset.delete()
    db.commit()
    return success_response(data={}, message="删除成功")


class GenerateImageResultCRUD(ModelCRUD):
    model = GenerateImageResult


@router.post("/del_generate_image_result", summary="删除生成图片", response_model=generate_image_schema.GenerateImageRecordSchemaRes, name='删除生成图片')
async def del_generate_image(
    *,
    db: Session = Depends(get_db),
    req_data: generate_image_schema.GenerateImageRecordSchemaDel,
    current_user: Users = Depends(security.verify_admin_permission),
) -> Any:
    crud = GenerateImageResultCRUD()
    queryset = crud.get_all_queryset(db)
    req_data = jsonable_encoder(req_data)
    queryset = queryset.filter(crud.model.id.in_(req_data.get('id', [])))
    for instance in queryset:
        delete_storage_file(instance.path)
    queryset.delete()
    db.commit()
    return success_response(data={}, message="删除成功")


